package com.ai.module.mj.service.impl;

import cn.hutool.core.bean.BeanUtil;
import cn.hutool.core.text.CharSequenceUtil;
import cn.hutool.core.util.RandomUtil;
import com.ai.module.chat.common.base.ThreadLocalContext;
import com.ai.module.chat.common.constant.ChatImagePointStringConstant;
import com.ai.module.chat.common.constant.StringConstant;
import com.ai.module.chat.common.domain.bo.AliUploadBo;
import com.ai.module.chat.common.enu.*;
import com.ai.module.chat.config.advice.exception.BusinessException;
import com.ai.module.business.domain.entity.BannerRecordEntity;
import com.ai.module.business.enu.Source;
import com.ai.module.business.service.BannerRecordService;
import com.ai.module.business.service.MemberCenterService;
import com.ai.module.mj.constant.enums.TaskStatus;
import com.ai.module.mj.constant.stringConstant.Constants;
import com.ai.module.mj.constant.stringConstant.ReturnCode;
import com.ai.module.mj.domain.Task;
import com.ai.module.mj.domain.bo.MJDataUrl;
import com.ai.module.mj.domain.bo.MJDataUrlSerializer;
import com.ai.module.mj.domain.bo.UserPointsBo;
import com.ai.module.mj.domain.dto.TaskCondition;
import com.ai.module.mj.domain.enu.ImageTaskSource;
import com.ai.module.mj.domain.req.*;
import com.ai.module.mj.domain.vo.SubmitResultVO;
import com.ai.module.mj.domain.vo.TaskChangeParams;
import com.ai.module.mj.service.HighTaskStoreService;
import com.ai.module.mj.service.ImageSubmitService;
import com.ai.module.user.domain.vo.UserLoginInfo;
import com.ai.module.user.enu.PointSource;
import com.ai.module.chat.util.Base64Util;
import com.ai.module.chat.util.DateUtils;
import com.ai.module.chat.util.aliyun.AliyunUtil;
import com.ai.module.mj.config.propertites.YmlProperties;
import com.ai.module.mj.constant.enums.TaskAction;
import com.ai.module.mj.service.TaskService;
import com.ai.module.mj.service.TranslateService;
import com.ai.module.mj.util.ConvertUtils;
import com.ai.module.mj.util.MimeTypeUtils;
import com.ai.module.mj.util.SnowFlakeUtil;
import com.aliyun.green20220302.models.ImageModerationResponseBody;
import lombok.extern.slf4j.Slf4j;
import org.jetbrains.annotations.NotNull;
import org.springframework.beans.factory.annotation.Autowired;
import org.springframework.stereotype.Service;
import org.springframework.util.CollectionUtils;
import org.springframework.util.StringUtils;

import java.io.File;
import java.net.MalformedURLException;
import java.util.*;

import static com.ai.module.chat.common.constant.FileStringConstant.MID_JOURNEY_IMAGE;
import static com.ai.module.chat.common.constant.FileStringConstant.MID_JOURNEY_USER;

/**
 * @author: MrHe
 * @create: 2024-01-06 17:22
 * @description:
 **/
@Service
@Slf4j
public class ImageSubmitServiceImpl implements ImageSubmitService {

    @Autowired
    private TranslateService translateService;
    @Autowired
    private HighTaskStoreService taskStoreService;
    @Autowired
    private YmlProperties properties;
    @Autowired
    private TaskService taskService;
    @Autowired
    private MemberCenterService memberCenterService;

    @Autowired
    private BannerRecordService bannerRecordService;

    private Boolean endFlag = Boolean.FALSE;

    @Override
    public void setEnd() {
        this.endFlag = true;
    }

    @Override
    public SubmitResultVO imagine(SubmitImagineReq imagineDTO) {
        UserLoginInfo userLoginInfo = ThreadLocalContext.getUserLoginInfo();
        checkPoints(imagineDTO, userLoginInfo);
        String prompt = imagineDTO.getPrompt();
        ImageTaskSource taskSource = imagineDTO.getTaskSource();
        if (CharSequenceUtil.isBlank(prompt)
                && taskSource == ImageTaskSource.commonImage
        ) {
            throw new BusinessException(ReturnCode.VALIDATION_ERROR, "prompt不能为空");
        }
        prompt = prompt.trim();
        String userId = ThreadLocalContext.getUserId();
        List<String> originalUrlList = imagineDTO.getFileUrlList();
        String promptEn = translateService.commonTranslatePromptEn(prompt);

        List<String> base64Array = Optional.ofNullable(imagineDTO.getBase64Array()).orElse(new ArrayList<>());
        if (CharSequenceUtil.isNotBlank(imagineDTO.getBase64())) {
            base64Array.add(imagineDTO.getBase64());
        }
        Task task = newTask(imagineDTO);
        List<MJDataUrl> dataUrls;
        try {
            if (CollectionUtils.isEmpty(originalUrlList)){
                dataUrls = ConvertUtils.convertBase64Array(base64Array);
            }else {
//                List<String> fileUrlList = dealFileUrlList(base64Array,
//                        userLoginInfo.getId(),
//                        TaskAction.IMAGINE
//                );
                imagineDTO.setFileUrlList(originalUrlList);
                task.setFileUrlList(originalUrlList);
            }
        } catch (MalformedURLException e) {
            throw new BusinessException(ReturnCode.VALIDATION_ERROR, "base64格式错误");
        }
        task.setAction(TaskAction.IMAGINE);
        task.setPrompt(prompt);
        task.setUserId(userId);
        task.setImageGroundPrompt(prompt);
        task.setUserIntId(ThreadLocalContext.getUserLoginInfo().getId());
        //task.setFileUrlList(base64Array);

        task.setPromptEn(promptEn);
        task.setDescription("/imagine " + prompt);
//        task.setDataUrls(dataUrls);
        return this.taskService.submitImagine(task);

    }

    private  void checkPoints(SubmitImagineReq imagineDTO, UserLoginInfo userLoginInfo) {
        String speedType = imagineDTO.getSpeedType();
        if (StringConstant.FAST.equalsIgnoreCase(speedType)) {
            // 快速模型
            memberCenterService.checkUserPoints(userLoginInfo, ChatImagePointStringConstant.fastCost);

        } else {
            memberCenterService.checkUserPoints(userLoginInfo, ChatImagePointStringConstant.slowCost);
        }
    }

    @NotNull
    public  List<String> dealFileUrlList(List<String> base64Array,
                                         int userId,
                                         TaskAction taskAction ) {
        AliUploadBo aliUploadBo = new AliUploadBo();
        List<String> fileUrlList = new ArrayList<>();
        for (String base64 : base64Array) {
            String projectName = String.format("%s/%s/%s/%s/%s",
                    MID_JOURNEY_IMAGE, MID_JOURNEY_USER,
                    userId,
                    taskAction,
                    DateUtils.getTransDateTime());
            aliUploadBo.setProjectName(projectName);
            String realBase64 = base64
                    .replaceAll("data:image/png;base64,", "")
                    .replaceAll("data:image/jpeg;base64,", "")
                    ;
            File file = Base64Util.base2Image(realBase64, "jpg");
            try {
                String aliFileUrl = AliyunUtil.uploadOSSFile(file, aliUploadBo);
                ImageModerationResponseBody.ImageModerationResponseBodyDataResult result = AliyunUtil.doRealCheckImage(aliFileUrl);
                String bannerLabel = result.getLabel();
                // 图片审核
                if (StringUtils.hasText(bannerLabel)){
                    BannerRecordEntity bannerRecord = new BannerRecordEntity();
                    String desc = AliyunImageCheck.getDescByName(bannerLabel);
                    bannerRecord.setRule(desc);
                    bannerRecord.setConfidence(result.getConfidence());
                    bannerRecord.setBusinessId(null);
                    bannerRecord.setOriginalRule(bannerLabel);
                    bannerRecord.setBusinessContent(aliFileUrl);
                    bannerRecord.setUserId(ThreadLocalContext.getUserId());
                    bannerRecord.setSource(Source.USER_DRAW_IMAGE);
                    bannerRecordService.save(bannerRecord);
                    if (desc != null) {
                        throw new BusinessException(desc.split("。")[0], ReturnCode.FAILURE);
                    }
                }
                fileUrlList.add(aliFileUrl);
            } catch (Exception e) {
                e.printStackTrace();
                if (e instanceof BusinessException) {
                    throw new BusinessException(e.getMessage());
                }
                throw new RuntimeException(e);
            }


        }
        return fileUrlList;
    }



    @Override
    public SubmitResultVO simpleChange(SubmitSimpleChangeReq simpleChangeDTO) {

        TaskChangeParams changeParams = ConvertUtils.convertChangeParams(simpleChangeDTO.getContent());
        if (changeParams == null) {
            throw new BusinessException(ReturnCode.VALIDATION_ERROR, "content参数错误");
        }
      ;
        SubmitChangeReq changeDTO = new SubmitChangeReq();
        BeanUtil.copyProperties(changeParams, changeDTO);
        changeDTO.setAction(changeParams.getAction());
        changeDTO.setRobotId(simpleChangeDTO.getRobotId());
//        changeDTO.setRobotId(simpleChangeDTO.getRobotId());
//        changeDTO.setTaskId(changeParams.getTaskId());
//        changeDTO.setIndex(changeParams.getIndex());
//        changeDTO.setState(simpleChangeDTO.getState());
//        changeDTO.setNotifyHook(simpleChangeDTO.getNotifyHook());
//        changeDTO.setPrompt();
        return change(changeDTO);
    }

    private void dealChangePoints(SubmitChangeReq changeParams ) {
        UserPointsBo userPointsBo = new UserPointsBo();
        userPointsBo.setPrompt(changeParams.getPrompt());
        userPointsBo.setModelType(ModelType.MID_JOURNEY);
        userPointsBo.setPointType(changeParams.getAction());
        String speedType = changeParams.getSpeedType();
        UserLoginInfo userLoginInfo = ThreadLocalContext.getUserLoginInfo();
        if (StringConstant.FAST.equalsIgnoreCase(speedType)) {
            // 快速模型
            memberCenterService.checkUserPoints(userLoginInfo, ChatImagePointStringConstant.fastCostUV);
        } else {
            memberCenterService.checkUserPoints(userLoginInfo, ChatImagePointStringConstant.slowCostUV);
        }
        userPointsBo.setBusinessId(changeParams.getUuid());
        userPointsBo.setSourceId(changeParams.getRobotId());
        userPointsBo.setSource(PointSource.commonImage);
        // memberCenterService.dealMidJourneyUserPoints(userPointsBo, task);
    }

    @Override
    public SubmitResultVO change(SubmitChangeReq changeDTO) {
        String description = "/up " + changeDTO.getTaskId();
        if (TaskAction.REROLL.equals(changeDTO.getAction())) {
            description += " R";
        } else {
            description += " " + changeDTO.getAction().name().charAt(0) + changeDTO.getIndex();
        }
        if (TaskAction.UPSCALE.equals(changeDTO.getAction())) {
            TaskCondition condition = new TaskCondition().setDescription(description);
            Task existTask = this.taskStoreService.findOne(condition);
            if (existTask != null) {
                return SubmitResultVO.of(ReturnCode.EXISTED, "任务已存在", existTask.getId())
                        .setProperty("status", existTask.getStatus())
                        .setProperty("imageUrl", existTask.getImageUrl());
            }
        }
        Task targetTask = this.taskStoreService.get(changeDTO.getTaskId());
        if (targetTask == null) {
            throw new BusinessException(ReturnCode.NOT_FOUND, "关联任务不存在或已失效");
        }
        if (!TaskStatus.SUCCESS.equals(targetTask.getStatus())) {
            throw new BusinessException(ReturnCode.VALIDATION_ERROR, "关联任务状态错误");
        }
        TaskAction action = targetTask.getAction();
        if (!Set.of(TaskAction.IMAGINE, TaskAction.VARIATION, TaskAction.REROLL, TaskAction.BLEND).contains(action)) {
            throw new BusinessException(ReturnCode.VALIDATION_ERROR, "关联任务不允许执行变化");
        }
        Task task = newTask(changeDTO);
        task.setAction(changeDTO.getAction());
        task.setPrompt(changeDTO.getPrompt());
        if (action== TaskAction.BLEND) {
            task.setImageGroundPrompt("源于用户拼图");
        }else {
            task.setImageGroundPrompt(targetTask.getPrompt());
        }
        task.setRealPromptEn(targetTask.getPromptEn());
        task.setFinalPrompt(targetTask.getProperty(Constants.TASK_PROPERTY_FINAL_PROMPT).toString());
        task.setProperty(Constants.TASK_PROPERTY_FINAL_PROMPT, targetTask.getProperty(Constants.TASK_PROPERTY_FINAL_PROMPT));
        task.setProperty(Constants.TASK_PROPERTY_PROGRESS_MESSAGE_ID, targetTask.getProperty(Constants.TASK_PROPERTY_MESSAGE_ID));
        task.setProperty(Constants.TASK_PROPERTY_DISCORD_INSTANCE_ID, targetTask.getProperty(Constants.TASK_PROPERTY_DISCORD_INSTANCE_ID));
        task.setDescription(description);

        int messageFlags = targetTask.getPropertyGeneric(Constants.TASK_PROPERTY_FLAGS);
        String messageId = targetTask.getPropertyGeneric(Constants.TASK_PROPERTY_MESSAGE_ID);
        String messageHash = targetTask.getPropertyGeneric(Constants.TASK_PROPERTY_MESSAGE_HASH);
        // 图片变换积分处理
        dealChangePoints(changeDTO);
        if (TaskAction.UPSCALE.equals(changeDTO.getAction())) {

            return this.taskService.submitUpscale(task, messageId, messageHash, changeDTO.getIndex(), messageFlags);
        } else if (TaskAction.VARIATION.equals(changeDTO.getAction())) {
            return this.taskService.submitVariation(task, messageId, messageHash, changeDTO.getIndex(), messageFlags);
        } else {
            return this.taskService.submitReroll(task, messageId, messageHash, messageFlags);
        }
    }

    @Override
    public SubmitResultVO describe(SubmitDescribeReq describeDTO) {
        if (CharSequenceUtil.isBlank(describeDTO.getBase64()) &&
                CollectionUtils.isEmpty(describeDTO .getFileUrlList())) {
            throw new BusinessException(ReturnCode.VALIDATION_ERROR, "base64或者图片地址不能为空");
        }
        MJDataUrlSerializer serializer = new MJDataUrlSerializer();
        MJDataUrl dataUrl= null;
        String taskFileName = null;
        List<String> fileUrlList = describeDTO.getFileUrlList();
        Task task = newTask(describeDTO);

        try {
            if (CollectionUtils.isEmpty(fileUrlList)){
                dataUrl = serializer.unserialize(describeDTO.getBase64());
                taskFileName = task.getId() + "." + MimeTypeUtils.guessFileSuffix(dataUrl.getMimeType());
            }else {
                taskFileName = fileUrlList.get(0);
            }
        } catch (MalformedURLException e) {
            throw new BusinessException(ReturnCode.VALIDATION_ERROR, "base64格式错误");
        }
        task.setAction(TaskAction.DESCRIBE);
        task.setDescription("/describe " + taskFileName);
        return this.taskService.submitDescribe(task, dataUrl,fileUrlList);

    }

    @Override
    public SubmitResultVO blend(SubmitBlendReq blendDTO) {
        List<String> base64Array = blendDTO.getBase64Array();
        List<String> realFileList = blendDTO.getFileUrlList();
        if (blendDTO.getDimensions() == null) {
            throw new BusinessException(ReturnCode.VALIDATION_ERROR, "dimensions参数错误");
        }
        MJDataUrlSerializer serializer = new MJDataUrlSerializer();
        List<MJDataUrl> dataUrlList = new ArrayList<>();
        List<String> fileUrlList = realFileList;
        try {

            for (String url : fileUrlList) {
                String base64ByUrl = Base64Util.getBase64ByUrl(url);
                base64Array.add(base64ByUrl);
            }
            realFileList= new ArrayList<>();
            if (CollectionUtils.isEmpty(realFileList)){
                fileUrlList = dealFileUrlList(base64Array,
                        ThreadLocalContext.getUserLoginInfo().getId(),
                        TaskAction.BLEND
                );
                for (String base64 : base64Array) {
                    MJDataUrl dataUrl = serializer.unserialize(base64);
                    dataUrlList.add(dataUrl);
                }
            }

        } catch (MalformedURLException e) {
            throw new BusinessException(ReturnCode.VALIDATION_ERROR, "base64格式错误");
        }

        blendDTO.setFileUrlList(fileUrlList);
        Task task = newTask(blendDTO);
        task.setAction(TaskAction.BLEND);
        task.setFileUrlList(fileUrlList);
        task.setDescription("/blend " + task.getId() + " " + dataUrlList.size());
        return this.taskService.submitBlend(task, dataUrlList, blendDTO.getDimensions());

    }

    private Task newTask(BaseSubmitReq base) {
        Task task = new Task();
        String prompt = base.getPrompt();
        task.setTaskSource(base.getTaskSource());
        task.setRobotId(base.getRobotId());
        task.setUserId(ThreadLocalContext.getUserId());
        task.setUserIntId(ThreadLocalContext.getUserLoginInfo().getId());
        task.setId(System.currentTimeMillis() + RandomUtil.randomNumbers(3));
        task.setSubmitTime(DateUtils.getCurrentTime());
        task.setState(base.getState());
        task.setUuid(base.getUuid());
        task.setLoginToken(ThreadLocalContext.getUserLoginInfo().getLoginToken());
        String notifyHook = CharSequenceUtil.isBlank(base.getNotifyHook()) ? this.properties.getNotifyHook() : base.getNotifyHook();
        task.setProperty(Constants.TASK_PROPERTY_NOTIFY_HOOK, notifyHook);
        task.setProperty(Constants.TASK_PROPERTY_NONCE, SnowFlakeUtil.INSTANCE.nextId());
        // 处理自定义参数
        MidJourneySpeed speed = base.getSpeed();
        StyleType styleType = base.getStyleType();
        Map<MidJourneyParam, String> midJourneyParams = base.getMidJourneyParams();
        StringBuilder sb = new StringBuilder();
        sb.append(speed.getSpeed());
        sb.append(styleType.getName());
        Set<MidJourneyParam> midJourneySet = midJourneyParams.keySet();
        for (MidJourneyParam journeyKey : midJourneySet) {
            // add 24.03.16 特殊情况 人像替换
            if (journeyKey == MidJourneyParam.CREF){
                String name = journeyKey.getName();
                StringBuilder crefSB = new StringBuilder();
                crefSB.append(name);
                crefSB.append("%s");
                String mjKey = journeyKey.getValue();
                String realValue = midJourneyParams.get(journeyKey);
                crefSB.append(mjKey);
                crefSB.append(realValue);
                task.setCrefPrompt(crefSB.toString());
                continue;
            }
            // 判断 是不是包含
            String mjKey = journeyKey.getValue();
            String realValue = midJourneyParams.get(journeyKey);
            if (StringUtils.hasText(prompt)
                    && !prompt.contains(realValue.replace(" ",""))
            ){
                sb.append(mjKey);
                sb.append(realValue);
            }
        }
        task.setOtherPrompt(sb.toString());
        return task;
    }



}
